package partition;

import earlyDetection.EarlyDetector;
import eval.Batch;
import network.Device;
import network.Network;
import network.Port;
import network.Rule;
import networkLoader.FBNetwork;
import util.Util;
import verifier.util.BDDEngine;
import verifier.util.PPM;
import verifier.JiffyVerifier;

import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.stream.Collectors;

public class SubspacePartition extends Thread {
    private Network network;
    private List<SubgraphPartition> subgraphPartitions = new ArrayList<>();
    private BDDEngine bddEngine;
    private boolean enableEarlyDetection;
    private BlockingQueue<Rule> pendingBatches = new LinkedBlockingQueue<>();
    private EarlyDetector earlyDetector = new EarlyDetector();

    public SubspacePartition(Network network, List<Rule> rules, boolean subgraph, boolean enableEarlyDetection) {
        this.enableEarlyDetection = enableEarlyDetection;
        this.network = network;
        this.bddEngine = new BDDEngine(32);
        this.subgraphPartitions = new ArrayList<>();
        if (subgraph) {
            HashMap<String, Partition> nameToPartition = Partitioner.setupPartitions(network, 112);
            for (Partition partition : nameToPartition.values()) {
                HashSet<Device> devInPartition = new HashSet<>(partition.getNetwork().getAllDevices());
                List<Rule> rulesInSubgraph = rules.stream().filter(rule -> devInPartition.contains(rule.getDevice())).collect(Collectors.toList());
                SubgraphPartition subgraphPartition = new SubgraphPartition(network, partition.getNetwork(), rulesInSubgraph, bddEngine);
                this.subgraphPartitions.add(subgraphPartition);
            }
        } else {
            SubgraphPartition subgraphPartition = new SubgraphPartition(network, network, rules, bddEngine);
            this.subgraphPartitions.add(subgraphPartition);
        }
    }

    public SubspacePartition(Network network, List<Rule> rules, boolean subgraph) {
        this(network, rules, subgraph, false);
    }

    public SubspacePartition(Network network, boolean enableSubgraph) {
        this.network = network;
        this.bddEngine = new BDDEngine(32);
        if (enableSubgraph) {
//            HashMap<String, Network> nameToSubgraphPartiton = FattreeNetwork.getPartitions(network, 32, this.bddEngine);
            HashMap<String, Network> nameToSubgraphPartiton = FBNetwork.getPartitions(network, 112);
            for (Network subgraphNetwork : nameToSubgraphPartiton.values()) {
                SubgraphPartition subgraphPartition = new SubgraphPartition(network, subgraphNetwork, null, bddEngine);
                this.subgraphPartitions.add(subgraphPartition);
            }
        } else {
            SubgraphPartition subgraphPartition = new SubgraphPartition(network, network, null, bddEngine);
            this.subgraphPartitions = new ArrayList<>();
            this.subgraphPartitions.add(subgraphPartition);
        }
    }

    private int nClosed = 0;
    private int nRules = 0;

    public void run() {
        while (true) {
            try {
                List<Rule> updates = new ArrayList<>();
                updates.add(pendingBatches.take());
                pendingBatches.drainTo(updates);
//                System.out.println("get batch with #rules: " + batch.getRules().size() + " #closed: " + batch.getClosed().size());
                if (this.subgraphPartitions.size() > 1) {
                    for (SubgraphPartition partition : this.subgraphPartitions) {
                        HashSet<Device> devInPartition = partition.getSubgraphNetwork().getAllDevices();
//                        List<Rule> rulesInSubgraph = batch.getRules().stream().filter(rule -> devInPartition.contains(rule.getDevice())).collect(Collectors.toList());
                        List<Rule> rulesInSubgraph = updates.stream().filter(rule -> devInPartition.contains(rule.getDevice())).collect(Collectors.toList());
                        if (rulesInSubgraph.size() > 0) {
                            partition.run(rulesInSubgraph);
                        }
                    }
                } else {
//                    subgraphPartitions.get(0).run(batch.getRules());
                    subgraphPartitions.get(0).run(updates);
                }

//                Dispatcher.processedRules += batch.getRules().size();

//                if (enableEarlyDetection) {
//                    // this call will start another thread do early detection
//                    earlyDetector.detectLoop(batch.getClosed(), getNetworkModel());
//                }
//                nClosed += batch.getClosed().size();
////                System.out.println(nClosed);
//                nRules += batch.getRules().size();
////                System.out.println(nRules);
//                if (nClosed == network.getAllDevices().size()) {
//                    Util.checkLoop(network.getAllDevices(), getNetworkModel());
//                }
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    public void run(List<Rule> rules, boolean enableFFMT) {
        if (this.subgraphPartitions.size() > 1) {
            for (SubgraphPartition partition : this.subgraphPartitions) {
                HashSet<Device> devInPartition = partition.getSubgraphNetwork().getAllDevices();
                List<Rule> rulesInSubgraph = rules.stream().filter(rule -> devInPartition.contains(rule.getDevice())).collect(Collectors.toList());
                if (rulesInSubgraph.size() > 0) {
                    long s = System.nanoTime();
                    partition.run(rulesInSubgraph, enableFFMT);
                    System.out.println("$model: " + (System.nanoTime() - s));
                    System.out.println("#rules: " + rulesInSubgraph.size());
                    System.out.println("#ECs: " + partition.getVerifier().predSize());
                }
            }
            Aggregator aggregator = new Aggregator(subgraphPartitions);
            long s = System.nanoTime();
            HashMap<Port, HashSet<Integer>> nm = aggregator.getNetworkModel();
            System.out.println("$global model: " + (System.nanoTime() - s));
        } else {
            long s = System.nanoTime();
            subgraphPartitions.get(0).run(rules, enableFFMT);
            System.out.println("$model: " + (System.nanoTime() - s));
            System.out.println("#rules: " + rules.size());
            System.out.println("#ECs: " + subgraphPartitions.get(0).getVerifier().predSize());
        }
    }

    public HashMap<Port, HashSet<Integer>> getNetworkModel() {
        if (subgraphPartitions.size() > 1) {
            Aggregator aggregator = new Aggregator(subgraphPartitions);
            long s = System.nanoTime();
            HashMap<Port, HashSet<Integer>> nm = aggregator.getNetworkModel();
            System.out.println("$global model: " + (System.nanoTime() - s));
            return nm;
        } else {
//            System.out.println("#ECs after batch: " + subgraphPartitions.get(0).getVerifier().portsToPredicate.size());
            JiffyVerifier verifier = subgraphPartitions.get(0).getVerifier();
            PPM ppm = new PPM(verifier.portsToPredicate);
            return ppm.getPortToPreds();
        }
    }

    public void run(boolean useFFMT) {
        for (SubgraphPartition subgraphPartition : this.subgraphPartitions) {
            subgraphPartition.run(useFFMT);
            System.gc();
        }
        if (subgraphPartitions.size() > 1) {
            Aggregator aggregator = new Aggregator(subgraphPartitions);
            long s = System.nanoTime();
            HashMap<Port, HashSet<Integer>> nm = aggregator.getNetworkModel();
            System.out.println("$global model: " + (System.nanoTime() - s));
        }
    }

    public void runEarly() {
        EarlyDetector earlyDetector = new EarlyDetector();
        HashSet<Device> closedDevices = new HashSet<>();
        for (SubgraphPartition subgraphPartition : this.subgraphPartitions) {
            Map<Device, List<Rule>> deviceToRules = new HashMap<>();
            for (Rule rule : subgraphPartition.getRules()) {
                deviceToRules.computeIfAbsent(rule.getDevice(), k -> new ArrayList<>());
                deviceToRules.get(rule.getDevice()).add(rule);
            }
            for (Map.Entry<Device, List<Rule>> entry : deviceToRules.entrySet()) {
                subgraphPartition.run(entry.getValue());
                closedDevices.add(entry.getKey());
                Aggregator aggregator = new Aggregator(this.subgraphPartitions);
                HashMap<Port, HashSet<Integer>> nm = aggregator.getNetworkModel();
//                earlyDetector.detectLoop(closedDevices, nm);
                if (earlyDetector.hasLoop()) {
                    return;
                }
            }

        }
    }

    public BlockingQueue<Rule> getPendingBatches() {
        return pendingBatches;
    }
}